
; (c) Daniel Llorens - 2012-2015
; No local dependencies beyond (assert).

; This library is free software; you can redistribute it and/or modify it under
; the terms of the GNU General Public License as published by the Free
; Software Foundation; either version 3 of the License, or (at your option) any
; later version.

(define-module (ploy basic))
(import (srfi srfi-1) (srfi srfi-8) (ice-9 format) (only (ploy assert) assert))

; @todo Array gradeup/down and dyadic (grade a by the order of b).
; also serves as permutation inverse.
; cf http://www.jsoftware.com/jwiki/Essays/Inverse%20Permutation
(define (gradeup l)
  (map cadr (sort (zip l (iota (length l)))
                  (lambda (a b) (< (car a) (car b))))))

(export gradeup)

; @todo do not use map.
; @todo do not run over the full lists for (length=? number . lists).
(define (length=? . a)
  (cond ((null? a) #t)
        ((number? (car a))
         (apply = (car a) (map length (cdr a))))
        ((null? (car a))
         (every null? (cdr a)))
        (else
         (and (not (any null? (cdr a)))
              (apply length=? (map cdr a))))))

(define (not= x y) (not (= x y)))

(export length=? not=)

(define array-fold
  (case-lambda
   ((op init)
    init)
   ((op init a0)
    (array-for-each (lambda (x0) (set! init (op x0 init))) a0)
    init)
   ((op init a0 a1)
    (array-for-each (lambda (x0 x1) (set! init (op x0 x1 init))) a0 a1)
    init)
   ((op init . args)
    (apply array-for-each (lambda x (set! init (apply op (append! x (list init))))) args)
    init)))

(export array-fold)

(define tally array-length)
(define ($ A) (if (array? A) (array-dimensions A) '()))
(define array-dim (lambda (A i) (list-ref ($ A) i)))
(define $. array-dim)

(define (make-array-of-size c type)
  (apply make-typed-array type *unspecified* ($ c)))

(define* (make-random-array dimensions #:key (type 'f64) (f values))
  (let ((a (apply make-typed-array type *unspecified* dimensions)))
    (array-map! (array-contents a) (lambda x (f (random:uniform))))
    a))

(define array-copy
  (case-lambda
    ((c) (array-copy #f c))
    ((type c)
     (let ((d (make-array-of-size c (or type (array-type c)))))
       (array-copy! c d)
       d))))

(define (array-map type op a0 . args)
  (if (array? a0)
    (let ((d (make-array-of-size a0 (or type (array-type a0)))))
      (apply array-map! d op a0 args)
      d)
    (apply op a0 args)))

(export tally $ array-dim $. make-array-of-size make-random-array
        array-copy array-map)

(define (index-rect s i) (fold (lambda (ss ii c) (+ ii (* ss c))) 0 s i))

(define (index-rect-lsd-first s i)
  (index-rect (reverse s) (reverse i)))

(export index-rect index-rect-lsd-first)

; This is used (redefined in tm's files) in box/jast/.geos.
(define (list-product . rest)
  "make a list of all lists with 1st element from 1st arg, 2nd element from
   2nd arg, etc.

   (list-product '(1 2)) => ((1) (2))
   (list-product '(1 2) '(3 4)) => ((1 3) (1 4) (2 3) (2 4))
   (list-product '(1 2) '(3 4) '(5 6)) => ((1 3 5) (1 3 6) etc.)"
  (cond
    ((null? rest)
      '())
    ((null? (cdr rest))
      (map list (car rest)))
    (else
      (let ((list-product-rest (apply list-product (cdr rest))))
        (append-map!
          (lambda (p)
            (map
              (lambda (v) (cons p v))
              list-product-rest))
          (car rest))))))

(export list-product)

; @todo Inefficient.
(define (interleave la lb)
  "Interleave two lists"
  (if (null? la)
    '()
    (append
      (list (car la) (car lb))
      (interleave (cdr la) (cdr lb)))))

(export interleave)

(define (list-accumulate* l op)
  (if (null? l)
    l
    (fold (lambda (a new) (cons (+ (op a) (car new)) new))
          (list (op (car l)))
          (cdr l))))

(define (list-accumulate l) (reverse! (list-accumulate* l values)))

(export list-accumulate* list-accumulate)

(define-syntax time
  (syntax-rules ()
    ((_ e0 ...)
      (let ((start (get-internal-run-time)))
        e0 ...
        (exact->inexact (/ (- (get-internal-run-time) start)
                           internal-time-units-per-second))))))

(define-syntax walltime
  (syntax-rules ()
    ((_ e0 ...)
      (let ((start (get-internal-real-time)))
        e0 ...
        (exact->inexact (/ (- (get-internal-real-time) start)
                           internal-time-units-per-second))))))

(define (UTC-date-string)
  (strftime "%Y-%m-%dT%H:%M:%SZ" (gmtime (current-time))))

(define (UTC-date-string-basic)
  (strftime "%Y%m%dT%H%M%SZ" (gmtime (current-time))))

(export time walltime UTC-date-string UTC-date-string-basic)

(define (format! s . args)
  (apply format #t (string-append s "\n~!") args))

(export format!)

(define (pk-shape s A) (format! "~a [~a]" s ($ A)) A)
(define (pk=> s f A) (format! "~a [~a]" s (f A)) A)

(export pk-shape pk=>)

(define (conj x) (make-rectangular (real-part x) (- (imag-part x))))
(define-inlinable (sqr x) (* x x))
(define-inlinable (sqrm x) (+ (sqr (real-part x)) (sqr (imag-part x))))

(export conj sqr sqrm)

; @todo replace with (over) when it's not slower.
(define sqrm-reduce
  (case-lambda
    ((a)
     (let ((end (tally a)))
       (let loop ((i 0) (c 0))
         (if (< i end)
           (loop (+ i 1) (+ c (sqrm (array-ref a i))))
           c))))
    ((a b)
     (let ((end (tally a)))
       (let loop ((i 0) (c 0))
         (if (< i end)
           (loop (+ i 1) (+ c (sqrm (- (array-ref a i) (array-ref b i)))))
           c))))))

(define sqrm-reduce*
  (case-lambda
    ((a)
     (array-fold (lambda (a c) (+ c (sqrm a))) 0 a))
    ((a b)
     (array-fold (lambda (a b c) (+ c (sqrm (- a b)))) 0 a b))))

;; @TODO sqrm-reduce* shouldn't be slower than sqrm-reduce.
;; (define A (array-copy 'f64 (i. #e1e7)))
;; (define B (array-copy #t (array-copy 'f64 (i. #e1e7))))
;; ,profile (sqrm-reduce A)

(define v-norm
  (case-lambda
    "v-norm a [b] => (sqrt (sqrm-reduce a [b]))"
    ((a) (sqrt (sqrm-reduce a)))
    ((a b) (sqrt (sqrm-reduce a b)))))

(export sqrm-reduce v-norm)

(define (rank A)
  (if (array? A) (array-rank A) 0))

(export rank)

; ---------------------------------------------------------------------
; reference implementations of functions in lloda-array-support branch in Guile.
; Uncomment if your Guile doesn't have them (untested).
; ---------------------------------------------------------------------

;; (define (array-slice a . i)
;;   (let ((li (length i)))
;;     (apply make-shared-array a
;;                (lambda t (append i t))
;;                (drop ($ a) li))))

;; (define (array-cell-ref a . i)
;;   (let ((li (length i)))
;;     (if (= (rank a) li)
;;         (apply array-ref a i)
;;         (apply make-shared-array a
;;                (lambda t (append i t))
;;                (drop ($ a) li)))))

;; (define (array-cell-set! a o . i)
;;   (let ((li (length i)))
;;     (if (= (rank a) li)
;;         (apply array-set! a o i)
;;         (array-copy! o (apply array-cell-ref a i)))))

;; (define (array-slice-for-each frank op . a)
;;   (let ((frame (take ($ (car a)) frank)))
;;     (unless (every (lambda (a) (equal? frame (take ($ a) frank))) (cdr a))
;;       (throw 'mismatched-argument-frames frank (map $ a)))
;;     (array-index-map! (apply make-shared-array (make-array #t) (const '()) frame)
;;                       (lambda i (apply op (map (lambda (a) (apply array-slice a i)) a))))))

;; (export array-slice array-cell-ref array-cell-set! array-slice-for-each)

; ---------------------------------------------------------------------
; misc
; ---------------------------------------------------------------------

; Slow, since array-index-map! is C and the λ can't be inlined.
(define (i. . args)
  (let ((a (apply make-array *unspecified* args)))
    (array-index-map! (array-contents a) identity)
    a))

(define iota.
  (case-lambda
    ((n z s)
     (if (negative? n)
       (error "bad n for iota." n)
       (let ((a (make-array *unspecified* n)))
         (array-index-map! a (lambda (i) (+ (* i s) z)))
         a)))
    ((n z) (iota. n z 1))
    ((n) (iota. n 0 1))))

(define (linspace. from to n)
  "(linspace. from to n) - Return vector of n elements between from and to"
  (cond ((<= n 1) (iota. n from))
        (else (iota. n from (/ (- to from) (- n 1))))))

(define (linspace-m. from to n)
  "(linspace-m. from to n) - Return vector of n elements between from and to"
  (cond ((<= n 1) #())
        (else (iota. (- n 1) from (/ (- to from) (- n 1))))))

(export i. iota. linspace. linspace-m.)

; @todo Actual type deduction.
(define (array-type* a)
  (cond ((array? a) (array-type a))
        ((and (number? a) (inexact? a))
         (if (real? a) 'f64 'c64))
        (else #t)))

(export array-type*)

; ---------------------------------------------------------------------
; reshape
; ---------------------------------------------------------------------

(define (ravel a)
  (or (array-contents a) (array-contents (array-copy (array-type* a) a))))

; Unravel, C order, repeat or clip, try to avoid copies. Equivalent to J (s $ ,a).
; Others:
; PDL reshape (pdl.perl.org/PDLdocs/Core.html#reshape) Unravel, pad with 0. Oddities.
; Octave reshape. Fortran order, requires size match.
; J dyad $ (http://www.jsoftware.com/help/jforc/declarations.htm#_Toc191734328). Does not unravel, works with items of a. Repeat or clip.
; Fortran RESHAPE (http://www.nsc.liu.se/~boein/f77to90/a5.html#section17). Unravel, Fortran order, requires size match unless pad is given, otherwise pad. Optional reorder.
(define (reshape a . s)
  "reshape A . s

   Reshape C-order ravel of array a to dimensions s, which must be nonnegative
   integers. One of the dimensions s may be #t, in which case it is deduced from
   the other dimensions and the size of A's ravel.

   Examples:

   (reshape (i. 2 3) 6) => #(0 1 2 3 4 5)
   (reshape (i. 2 3) 5) => #(0 1 2 3 4)
   (reshape (i. 2 3) 7) => #(0 1 2 3 4 5 0)
   (reshape (i. 2 3) 3 2) => #2((0 1) (2 3) (4 5))
   (reshape (i. 2 3) 2 2 2) => #3(((0 1) (2 3)) ((4 5) (0 1)))
   (reshape (i. 2 3) 4 2) => #2((0 1) (2 3) (4 5) (0 1))
   (reshape (i. 2 3) #t 2) => #2((0 1) (2 3) (4 5))
   (reshape (i. 2 3) 2 #t) => #2((0 1 2) (3 4 5))
   (reshape (i. 2 3) #t) => #(0 1 2 3 4 5)
   (reshape (i. 2 3) 0 3) => #2:0:3()
   (reshape (i. 2 3) 0) => #()
   (reshape (i. 2 3) 0 0 0) => #3()
   (reshape (i. 2 3) 0 #t) => error (can't deduce placeholder)
   (reshape (i. 2 3) 4 #t) => error (can't deduce placeholder)
   (reshape (i. 2 3)) => 0"

  (define mk make-shared-array)
  (if (array? a)
    (let* ((z ($ a)) (lz (apply * z)))
      (let loopb ((ss s) (ls 1))
; remove placeholder.
        (cond ((null? ss))
              ((boolean? (car ss))
               (let ((quot (apply * ls (cdr ss))))
                 (when (not (positive? quot)) (throw 'cannot-deduce-dimension-with-empty-result-shape quot))
                 (let ((pv (/ lz quot)))
                   (when (not (and (integer? pv) (>= pv 0))) (throw 'bad-placeholder pv))
                   (set-car! ss pv)
                   (set! ls lz))))
              ((integer? (car ss))
               (loopb (cdr ss) (* ls (car ss))))
              (else (throw 'bad-dimension-to-reshape (car ss))))
        (let loopr ((rz (reverse z)) (rs (reverse s)))
          (cond
; select.
           ((null? rs)
            (apply array-cell-ref a (make-list (length rz) 0)))
; tile.
           ((null? rz)
            (apply mk a (lambda i (drop i (length rs))) s))
           ((= (car rz) (car rs))
            (loopr (cdr rz) (cdr rs)))
           (else
            (let ((ls (apply * s))
                  (a1 (ravel a)))
              (if (>= lz ls)
; select ravel.
                (apply mk a1 (lambda i (list (index-rect s i))) s)
                (receive (q r) (floor/ ls lz)
; drop prefix strides.
                  (or (and (zero? r)
                           (let loops ((ss s) (p 1) (w 0))
                             (cond
                              ((= p q)
                               (apply mk a1 (lambda i (list (index-rect ss (drop i w)))) s))
                              ((< p q) (loops (cdr ss) (* p (car ss)) (+ 1 w)))
                              (else #f))))
; copy.
                      (let* ((b (apply make-typed-array (array-type* a) *unspecified* s))
                             (b1 (array-contents b)))
                        (array-copy!
                         (mk a1 (lambda i (cdr i)) q lz)
                         (mk b1 (lambda (i j) (list (+ j (* lz i)))) q lz))
                        (array-copy!
                         (mk a1 list r)
                         (mk b1 (lambda (j) (list (+ j (* lz q)))) r))
                        b))))))))))
    (apply reshape (make-array a) s)))

(define (reshape-strict A . s)
  (unless (= (apply * ($ A)) (apply * s))
    (throw 'mismatched-sizes ($ A) s))
  (apply reshape A s))

(export ravel reshape reshape-strict)
